import numpy as np
import pandas as pd
from scipy.stats import pearsonr
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import os
import multiprocessing
from tqdm import tqdm
import warnings
from bin_quant.backtest.Backtest import Backtest
from bin_quant.agent.quant_agent.FactorStrategy import FactorStrategy
warnings.filterwarnings('ignore')

def parallel_backtest(params, data):
    """并行执行的回测函数"""
    top_n, factor_weights = params
    strategy = FactorStrategy(data, factor_weights)
    backtest = Backtest(data, strategy, top_n=top_n)
    backtest.run()
    performance = backtest.evaluate_performance()
    return performance


def optimize_parameters(data, top_n_list=None, factor_combinations=None, n_jobs=-1):
    """
    优化策略参数
    :param data: 数据
    :param top_n_list: 要测试的top_n列表
    :param factor_combinations: 要测试的因子组合列表
    :param n_jobs: 并行工作数
    :return: 最佳参数和结果
    """
    if top_n_list is None:
        top_n_list = [5, 10, 15, 20]

    if factor_combinations is None:
        # 默认测试所有因子等权组合
        factor_cols = [col for col in data.columns if col not in
                       ['date', 'stock', 'open', 'high', 'low', 'close',
                        'volume', 'pe', 'pb', 'market_cap', 'return']]
        factor_combinations = [None]  # None表示等权

    # 生成参数组合
    param_combinations = [(top_n, weights) for top_n in top_n_list
                          for weights in factor_combinations]

    # 并行回测
    if n_jobs == -1:
        n_jobs = multiprocessing.cpu_count()

    with multiprocessing.Pool(n_jobs) as pool:
        results = list(tqdm(pool.imap(
            lambda params: parallel_backtest(params, data), param_combinations),
            total=len(param_combinations)))

    # 找到最佳参数
    best_result = max(results, key=lambda x: x['Sharpe Ratio'])
    best_idx = results.index(best_result)
    best_params = param_combinations[best_idx]

    return best_params, results
